Skip to content

[MemProf] Fix reporting with -memprof-matching-cold-threshold #173327

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
teresajohnson wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[MemProf] Fix reporting with -memprof-matching-cold-threshold #173327

teresajohnson wants to merge 1 commit into from
+60 −33

Conversation

@teresajohnson
Copy link
Contributor

@teresajohnson teresajohnson commented Dec 23, 2025

With the -memprof-matching-cold-threshold option, we hint as cold
allocations where the fraction of cold bytes is at least the given
threshold. However, we were incorrectly reporting all of the
allocation's contexts and bytes as hinted cold.

Fix this to report the non-cold contexts as ignored. To do this,
refactor out some existing reporting, and also keep track of the
original allocation type for each context in the Trie along with its
ContextTotalSize information. Most of the changes are the change to this
array's type and name.

@teresajohnson
[MemProf] Fix reporting with -memprof-matching-cold-threshold
b027b65 
With the -memprof-matching-cold-threshold option, we hint as cold
allocations where the fraction of cold bytes is at least the given
threshold. However, we were incorrectly reporting all of the
allocation's contexts and bytes as hinted cold.

Fix this to report the non-cold contexts as ignored. To do this,
refactor out some existing reporting, and also keep track of the
original allocation type for each context in the Trie along with its
ContextTotalSize information. Most of the changes are the change to this
array's type and name.
@llvmbot llvmbot added PGO Profile Guided Optimizations llvm:analysis Includes value tracking, cost tables and constant folding llvm:transforms labels Dec 23, 2025
@llvmbot
Copy link
Member

llvmbot commented Dec 23, 2025 •
edited
Loading

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-pgo

Author: Teresa Johnson (teresajohnson)

Changes

With the -memprof-matching-cold-threshold option, we hint as cold
allocations where the fraction of cold bytes is at least the given
threshold. However, we were incorrectly reporting all of the
allocation's contexts and bytes as hinted cold.

Fix this to report the non-cold contexts as ignored. To do this,
refactor out some existing reporting, and also keep track of the
original allocation type for each context in the Trie along with its
ContextTotalSize information. Most of the changes are the change to this
array's type and name.

Full diff: https://github.com/llvm/llvm-project/pull/173327.diff

3 Files Affected:

  • (modified) llvm/include/llvm/Analysis/MemoryProfileInfo.h (+9-6)
  • (modified) llvm/lib/Analysis/MemoryProfileInfo.cpp (+49-25)
  • (modified) llvm/test/Transforms/PGOProfile/memprof.ll (+2-2)
diff --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
index ba4010bd8f50c..dbc7a163bfd71 100644
--- a/llvm/include/llvm/Analysis/MemoryProfileInfo.h
+++ b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
@@ -75,8 +75,9 @@ class CallStackTrie {
     // If the user has requested reporting of hinted sizes, keep track of the
     // associated full stack id and profiled sizes. Can have more than one
     // after trimming (e.g. when building from metadata). This is only placed on
-    // the last (root-most) trie node for each allocation context.
-    std::vector<ContextTotalSize> ContextSizeInfo;
+    // the last (root-most) trie node for each allocation context. Also
+    // track the original allocation type of the context.
+    std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
     // Map of caller stack id to the corresponding child Trie node.
     std::map<uint64_t, CallStackTrieNode *> Callers;
     CallStackTrieNode(AllocationType Type)
@@ -118,10 +119,12 @@ class CallStackTrie {
     delete Node;
   }
 
-  // Recursively build up a complete list of context size information from the
-  // trie nodes reached form the given Node, for hint size reporting.
-  void collectContextSizeInfo(CallStackTrieNode *Node,
-                              std::vector<ContextTotalSize> &ContextSizeInfo);
+  // Recursively build up a complete list of context information from the
+  // trie nodes reached form the given Node, including each context's
+  // ContextTotalSize and AllocationType, for hint size reporting.
+  void collectContextInfo(
+      CallStackTrieNode *Node,
+      std::vector<std::pair<ContextTotalSize, AllocationType>> &ContextInfo);
 
   // Recursively convert hot allocation types to notcold, since we don't
   // actually do any cloning for hot contexts, to facilitate more aggressive
diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index b09f4ed78ca7e..bb1a103721740 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -181,7 +181,13 @@ void CallStackTrie::addCallStack(
     Curr = New;
   }
   assert(Curr);
-  llvm::append_range(Curr->ContextSizeInfo, ContextSizeInfo);
+  // Append all of the ContextSizeInfo, along with their original AllocType.
+  llvm::append_range(
+      Curr->ContextInfo,
+      llvm::map_range(
+          ContextSizeInfo, [AllocType](const ContextTotalSize &CTS) {
+            return std::pair<ContextTotalSize, AllocationType>(CTS, AllocType);
+          }));
 }
 
 void CallStackTrie::addCallStack(MDNode *MIB) {
@@ -214,18 +220,18 @@ void CallStackTrie::addCallStack(MDNode *MIB) {
   addCallStack(getMIBAllocType(MIB), CallStack, std::move(ContextSizeInfo));
 }
 
-static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
-                             AllocationType AllocType,
-                             ArrayRef<ContextTotalSize> ContextSizeInfo,
-                             const uint64_t MaxColdSize,
-                             bool BuiltFromExistingMetadata,
-                             uint64_t &TotalBytes, uint64_t &ColdBytes) {
+static MDNode *
+createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
+              AllocationType AllocType,
+              ArrayRef<std::pair<ContextTotalSize, AllocationType>> ContextInfo,
+              const uint64_t MaxColdSize, bool BuiltFromExistingMetadata,
+              uint64_t &TotalBytes, uint64_t &ColdBytes) {
   SmallVector<Metadata *> MIBPayload(
       {buildCallstackMetadata(MIBCallStack, Ctx)});
   MIBPayload.push_back(
       MDString::get(Ctx, getAllocTypeAttributeString(AllocType)));
 
-  if (ContextSizeInfo.empty()) {
+  if (ContextInfo.empty()) {
     // The profile matcher should have provided context size info if there was a
     // MinCallsiteColdBytePercent < 100. Here we check >=100 to gracefully
     // handle a user-provided percent larger than 100. However, we may not have
@@ -234,7 +240,8 @@ static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
     return MDNode::get(Ctx, MIBPayload);
   }
 
-  for (const auto &[FullStackId, TotalSize] : ContextSizeInfo) {
+  for (const auto &[CSI, AT] : ContextInfo) {
+    const auto &[FullStackId, TotalSize] = CSI;
     TotalBytes += TotalSize;
     bool LargeColdContext = false;
     if (AllocType == AllocationType::Cold) {
@@ -267,11 +274,12 @@ static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
   return MDNode::get(Ctx, MIBPayload);
 }
 
-void CallStackTrie::collectContextSizeInfo(
-    CallStackTrieNode *Node, std::vector<ContextTotalSize> &ContextSizeInfo) {
-  llvm::append_range(ContextSizeInfo, Node->ContextSizeInfo);
+void CallStackTrie::collectContextInfo(
+    CallStackTrieNode *Node,
+    std::vector<std::pair<ContextTotalSize, AllocationType>> &ContextInfo) {
+  llvm::append_range(ContextInfo, Node->ContextInfo);
   for (auto &Caller : Node->Callers)
-    collectContextSizeInfo(Caller.second, ContextSizeInfo);
+    collectContextInfo(Caller.second, ContextInfo);
 }
 
 void CallStackTrie::convertHotToNotCold(CallStackTrieNode *Node) {
@@ -283,6 +291,17 @@ void CallStackTrie::convertHotToNotCold(CallStackTrieNode *Node) {
     convertHotToNotCold(Caller.second);
 }
 
+// Helper to emit messages for non-cold contexts that are ignored for various
+// reasons when reporting of hinted bytes is enabled.
+static void emitIgnoredNonColdContextMessage(StringRef Tag,
+                                             uint64_t FullStackId,
+                                             StringRef Extra,
+                                             uint64_t TotalSize) {
+  errs() << "MemProf hinting: Total size for " << Tag
+         << " non-cold full allocation context hash " << FullStackId << Extra
+         << ": " << TotalSize << "\n";
+}
+
 // Copy over some or all of NewMIBNodes to the SavedMIBNodes vector, depending
 // on options that enable filtering out some NotCold contexts.
 static void saveFilteredNewMIBNodes(std::vector<Metadata *> &NewMIBNodes,
@@ -321,9 +340,7 @@ static void saveFilteredNewMIBNodes(std::vector<Metadata *> &NewMIBNodes,
       uint64_t TS =
           mdconst::dyn_extract<ConstantInt>(ContextSizePair->getOperand(1))
               ->getZExtValue();
-      errs() << "MemProf hinting: Total size for " << Tag
-             << " non-cold full allocation context hash " << FullStackId
-             << Extra << ": " << TS << "\n";
+      emitIgnoredNonColdContextMessage(Tag, FullStackId, Extra, TS);
     }
   };
 
@@ -430,10 +447,10 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
   // Trim context below the first node in a prefix with a single alloc type.
   // Add an MIB record for the current call stack prefix.
   if (hasSingleAllocType(Node->AllocTypes)) {
-    std::vector<ContextTotalSize> ContextSizeInfo;
-    collectContextSizeInfo(Node, ContextSizeInfo);
+    std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
+    collectContextInfo(Node, ContextInfo);
     MIBNodes.push_back(createMIBNode(
-        Ctx, MIBCallStack, (AllocationType)Node->AllocTypes, ContextSizeInfo,
+        Ctx, MIBCallStack, (AllocationType)Node->AllocTypes, ContextInfo,
         MaxColdSize, BuiltFromExistingMetadata, TotalBytes, ColdBytes));
     return true;
   }
@@ -486,10 +503,10 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
   // non-cold allocation type.
   if (!CalleeHasAmbiguousCallerContext)
     return false;
-  std::vector<ContextTotalSize> ContextSizeInfo;
-  collectContextSizeInfo(Node, ContextSizeInfo);
+  std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
+  collectContextInfo(Node, ContextInfo);
   MIBNodes.push_back(createMIBNode(
-      Ctx, MIBCallStack, AllocationType::NotCold, ContextSizeInfo, MaxColdSize,
+      Ctx, MIBCallStack, AllocationType::NotCold, ContextInfo, MaxColdSize,
       BuiltFromExistingMetadata, TotalBytes, ColdBytes));
   return true;
 }
@@ -503,9 +520,16 @@ void CallStackTrie::addSingleAllocTypeAttribute(CallBase *CI, AllocationType AT,
   removeAnyExistingAmbiguousAttribute(CI);
   CI->addFnAttr(A);
   if (MemProfReportHintedSizes) {
-    std::vector<ContextTotalSize> ContextSizeInfo;
-    collectContextSizeInfo(Alloc, ContextSizeInfo);
-    for (const auto &[FullStackId, TotalSize] : ContextSizeInfo) {
+    std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
+    collectContextInfo(Alloc, ContextInfo);
+    for (const auto &[CSI, OrigAT] : ContextInfo) {
+      const auto &[FullStackId, TotalSize] = CSI;
+      // If the original alloc type is not the one being applied as the hint,
+      // report that we ignored this context.
+      if (AT != OrigAT) {
+        emitIgnoredNonColdContextMessage("ignored", FullStackId, "", TotalSize);
+        continue;
+      }
       errs() << "MemProf hinting: Total size for full allocation context hash "
              << FullStackId << " and " << Descriptor << " alloc type "
              << getAllocTypeAttributeString(AT) << ": " << TotalSize << "\n";
diff --git a/llvm/test/Transforms/PGOProfile/memprof.ll b/llvm/test/Transforms/PGOProfile/memprof.ll
index f6a89a8ceb86a..a1f0f1d403c8f 100644
--- a/llvm/test/Transforms/PGOProfile/memprof.ll
+++ b/llvm/test/Transforms/PGOProfile/memprof.ll
@@ -400,10 +400,10 @@ for.end:                                          ; preds = %for.cond
 ;; with the full allocation context hash, type, and size in bytes.
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 8525406123785421946 and dominant alloc type cold: 10
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 11714230664165068698 and dominant alloc type cold: 10
-; TOTALSIZESTHRESH60: Total size for full allocation context hash 5725971306423925017 and dominant alloc type cold: 10
+; TOTALSIZESTHRESH60: Total size for ignored non-cold full allocation context hash 5725971306423925017: 10
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 16342802530253093571 and dominant alloc type cold: 10
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 18254812774972004394 and dominant alloc type cold: 10
-; TOTALSIZESTHRESH60: Total size for full allocation context hash 1093248920606587996 and dominant alloc type cold: 10
+; TOTALSIZESTHRESH60: Total size for ignored non-cold full allocation context hash 1093248920606587996: 10
 ; TOTALSIZESSINGLE: Total size for full allocation context hash 6792096022461663180 and single alloc type notcold: 10
 ; REMARKSINGLE: remark: memprof.cc:25:13: call in function main marked with memprof allocation attribute notcold
 ; TOTALSIZESSINGLE: Total size for full allocation context hash 15737101490731057601 and single alloc type cold: 10

@llvmbot
Copy link
Member

llvmbot commented Dec 23, 2025

@llvm/pr-subscribers-llvm-analysis

Author: Teresa Johnson (teresajohnson)

Changes

With the -memprof-matching-cold-threshold option, we hint as cold
allocations where the fraction of cold bytes is at least the given
threshold. However, we were incorrectly reporting all of the
allocation's contexts and bytes as hinted cold.

Fix this to report the non-cold contexts as ignored. To do this,
refactor out some existing reporting, and also keep track of the
original allocation type for each context in the Trie along with its
ContextTotalSize information. Most of the changes are the change to this
array's type and name.

Full diff: https://github.com/llvm/llvm-project/pull/173327.diff

3 Files Affected:

  • (modified) llvm/include/llvm/Analysis/MemoryProfileInfo.h (+9-6)
  • (modified) llvm/lib/Analysis/MemoryProfileInfo.cpp (+49-25)
  • (modified) llvm/test/Transforms/PGOProfile/memprof.ll (+2-2)
diff --git a/llvm/include/llvm/Analysis/MemoryProfileInfo.h b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
index ba4010bd8f50c..dbc7a163bfd71 100644
--- a/llvm/include/llvm/Analysis/MemoryProfileInfo.h
+++ b/llvm/include/llvm/Analysis/MemoryProfileInfo.h
@@ -75,8 +75,9 @@ class CallStackTrie {
     // If the user has requested reporting of hinted sizes, keep track of the
     // associated full stack id and profiled sizes. Can have more than one
     // after trimming (e.g. when building from metadata). This is only placed on
-    // the last (root-most) trie node for each allocation context.
-    std::vector<ContextTotalSize> ContextSizeInfo;
+    // the last (root-most) trie node for each allocation context. Also
+    // track the original allocation type of the context.
+    std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
     // Map of caller stack id to the corresponding child Trie node.
     std::map<uint64_t, CallStackTrieNode *> Callers;
     CallStackTrieNode(AllocationType Type)
@@ -118,10 +119,12 @@ class CallStackTrie {
     delete Node;
   }
 
-  // Recursively build up a complete list of context size information from the
-  // trie nodes reached form the given Node, for hint size reporting.
-  void collectContextSizeInfo(CallStackTrieNode *Node,
-                              std::vector<ContextTotalSize> &ContextSizeInfo);
+  // Recursively build up a complete list of context information from the
+  // trie nodes reached form the given Node, including each context's
+  // ContextTotalSize and AllocationType, for hint size reporting.
+  void collectContextInfo(
+      CallStackTrieNode *Node,
+      std::vector<std::pair<ContextTotalSize, AllocationType>> &ContextInfo);
 
   // Recursively convert hot allocation types to notcold, since we don't
   // actually do any cloning for hot contexts, to facilitate more aggressive
diff --git a/llvm/lib/Analysis/MemoryProfileInfo.cpp b/llvm/lib/Analysis/MemoryProfileInfo.cpp
index b09f4ed78ca7e..bb1a103721740 100644
--- a/llvm/lib/Analysis/MemoryProfileInfo.cpp
+++ b/llvm/lib/Analysis/MemoryProfileInfo.cpp
@@ -181,7 +181,13 @@ void CallStackTrie::addCallStack(
     Curr = New;
   }
   assert(Curr);
-  llvm::append_range(Curr->ContextSizeInfo, ContextSizeInfo);
+  // Append all of the ContextSizeInfo, along with their original AllocType.
+  llvm::append_range(
+      Curr->ContextInfo,
+      llvm::map_range(
+          ContextSizeInfo, [AllocType](const ContextTotalSize &CTS) {
+            return std::pair<ContextTotalSize, AllocationType>(CTS, AllocType);
+          }));
 }
 
 void CallStackTrie::addCallStack(MDNode *MIB) {
@@ -214,18 +220,18 @@ void CallStackTrie::addCallStack(MDNode *MIB) {
   addCallStack(getMIBAllocType(MIB), CallStack, std::move(ContextSizeInfo));
 }
 
-static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
-                             AllocationType AllocType,
-                             ArrayRef<ContextTotalSize> ContextSizeInfo,
-                             const uint64_t MaxColdSize,
-                             bool BuiltFromExistingMetadata,
-                             uint64_t &TotalBytes, uint64_t &ColdBytes) {
+static MDNode *
+createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
+              AllocationType AllocType,
+              ArrayRef<std::pair<ContextTotalSize, AllocationType>> ContextInfo,
+              const uint64_t MaxColdSize, bool BuiltFromExistingMetadata,
+              uint64_t &TotalBytes, uint64_t &ColdBytes) {
   SmallVector<Metadata *> MIBPayload(
       {buildCallstackMetadata(MIBCallStack, Ctx)});
   MIBPayload.push_back(
       MDString::get(Ctx, getAllocTypeAttributeString(AllocType)));
 
-  if (ContextSizeInfo.empty()) {
+  if (ContextInfo.empty()) {
     // The profile matcher should have provided context size info if there was a
     // MinCallsiteColdBytePercent < 100. Here we check >=100 to gracefully
     // handle a user-provided percent larger than 100. However, we may not have
@@ -234,7 +240,8 @@ static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
     return MDNode::get(Ctx, MIBPayload);
   }
 
-  for (const auto &[FullStackId, TotalSize] : ContextSizeInfo) {
+  for (const auto &[CSI, AT] : ContextInfo) {
+    const auto &[FullStackId, TotalSize] = CSI;
     TotalBytes += TotalSize;
     bool LargeColdContext = false;
     if (AllocType == AllocationType::Cold) {
@@ -267,11 +274,12 @@ static MDNode *createMIBNode(LLVMContext &Ctx, ArrayRef<uint64_t> MIBCallStack,
   return MDNode::get(Ctx, MIBPayload);
 }
 
-void CallStackTrie::collectContextSizeInfo(
-    CallStackTrieNode *Node, std::vector<ContextTotalSize> &ContextSizeInfo) {
-  llvm::append_range(ContextSizeInfo, Node->ContextSizeInfo);
+void CallStackTrie::collectContextInfo(
+    CallStackTrieNode *Node,
+    std::vector<std::pair<ContextTotalSize, AllocationType>> &ContextInfo) {
+  llvm::append_range(ContextInfo, Node->ContextInfo);
   for (auto &Caller : Node->Callers)
-    collectContextSizeInfo(Caller.second, ContextSizeInfo);
+    collectContextInfo(Caller.second, ContextInfo);
 }
 
 void CallStackTrie::convertHotToNotCold(CallStackTrieNode *Node) {
@@ -283,6 +291,17 @@ void CallStackTrie::convertHotToNotCold(CallStackTrieNode *Node) {
     convertHotToNotCold(Caller.second);
 }
 
+// Helper to emit messages for non-cold contexts that are ignored for various
+// reasons when reporting of hinted bytes is enabled.
+static void emitIgnoredNonColdContextMessage(StringRef Tag,
+                                             uint64_t FullStackId,
+                                             StringRef Extra,
+                                             uint64_t TotalSize) {
+  errs() << "MemProf hinting: Total size for " << Tag
+         << " non-cold full allocation context hash " << FullStackId << Extra
+         << ": " << TotalSize << "\n";
+}
+
 // Copy over some or all of NewMIBNodes to the SavedMIBNodes vector, depending
 // on options that enable filtering out some NotCold contexts.
 static void saveFilteredNewMIBNodes(std::vector<Metadata *> &NewMIBNodes,
@@ -321,9 +340,7 @@ static void saveFilteredNewMIBNodes(std::vector<Metadata *> &NewMIBNodes,
       uint64_t TS =
           mdconst::dyn_extract<ConstantInt>(ContextSizePair->getOperand(1))
               ->getZExtValue();
-      errs() << "MemProf hinting: Total size for " << Tag
-             << " non-cold full allocation context hash " << FullStackId
-             << Extra << ": " << TS << "\n";
+      emitIgnoredNonColdContextMessage(Tag, FullStackId, Extra, TS);
     }
   };
 
@@ -430,10 +447,10 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
   // Trim context below the first node in a prefix with a single alloc type.
   // Add an MIB record for the current call stack prefix.
   if (hasSingleAllocType(Node->AllocTypes)) {
-    std::vector<ContextTotalSize> ContextSizeInfo;
-    collectContextSizeInfo(Node, ContextSizeInfo);
+    std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
+    collectContextInfo(Node, ContextInfo);
     MIBNodes.push_back(createMIBNode(
-        Ctx, MIBCallStack, (AllocationType)Node->AllocTypes, ContextSizeInfo,
+        Ctx, MIBCallStack, (AllocationType)Node->AllocTypes, ContextInfo,
         MaxColdSize, BuiltFromExistingMetadata, TotalBytes, ColdBytes));
     return true;
   }
@@ -486,10 +503,10 @@ bool CallStackTrie::buildMIBNodes(CallStackTrieNode *Node, LLVMContext &Ctx,
   // non-cold allocation type.
   if (!CalleeHasAmbiguousCallerContext)
     return false;
-  std::vector<ContextTotalSize> ContextSizeInfo;
-  collectContextSizeInfo(Node, ContextSizeInfo);
+  std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
+  collectContextInfo(Node, ContextInfo);
   MIBNodes.push_back(createMIBNode(
-      Ctx, MIBCallStack, AllocationType::NotCold, ContextSizeInfo, MaxColdSize,
+      Ctx, MIBCallStack, AllocationType::NotCold, ContextInfo, MaxColdSize,
       BuiltFromExistingMetadata, TotalBytes, ColdBytes));
   return true;
 }
@@ -503,9 +520,16 @@ void CallStackTrie::addSingleAllocTypeAttribute(CallBase *CI, AllocationType AT,
   removeAnyExistingAmbiguousAttribute(CI);
   CI->addFnAttr(A);
   if (MemProfReportHintedSizes) {
-    std::vector<ContextTotalSize> ContextSizeInfo;
-    collectContextSizeInfo(Alloc, ContextSizeInfo);
-    for (const auto &[FullStackId, TotalSize] : ContextSizeInfo) {
+    std::vector<std::pair<ContextTotalSize, AllocationType>> ContextInfo;
+    collectContextInfo(Alloc, ContextInfo);
+    for (const auto &[CSI, OrigAT] : ContextInfo) {
+      const auto &[FullStackId, TotalSize] = CSI;
+      // If the original alloc type is not the one being applied as the hint,
+      // report that we ignored this context.
+      if (AT != OrigAT) {
+        emitIgnoredNonColdContextMessage("ignored", FullStackId, "", TotalSize);
+        continue;
+      }
       errs() << "MemProf hinting: Total size for full allocation context hash "
              << FullStackId << " and " << Descriptor << " alloc type "
              << getAllocTypeAttributeString(AT) << ": " << TotalSize << "\n";
diff --git a/llvm/test/Transforms/PGOProfile/memprof.ll b/llvm/test/Transforms/PGOProfile/memprof.ll
index f6a89a8ceb86a..a1f0f1d403c8f 100644
--- a/llvm/test/Transforms/PGOProfile/memprof.ll
+++ b/llvm/test/Transforms/PGOProfile/memprof.ll
@@ -400,10 +400,10 @@ for.end:                                          ; preds = %for.cond
 ;; with the full allocation context hash, type, and size in bytes.
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 8525406123785421946 and dominant alloc type cold: 10
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 11714230664165068698 and dominant alloc type cold: 10
-; TOTALSIZESTHRESH60: Total size for full allocation context hash 5725971306423925017 and dominant alloc type cold: 10
+; TOTALSIZESTHRESH60: Total size for ignored non-cold full allocation context hash 5725971306423925017: 10
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 16342802530253093571 and dominant alloc type cold: 10
 ; TOTALSIZESTHRESH60: Total size for full allocation context hash 18254812774972004394 and dominant alloc type cold: 10
-; TOTALSIZESTHRESH60: Total size for full allocation context hash 1093248920606587996 and dominant alloc type cold: 10
+; TOTALSIZESTHRESH60: Total size for ignored non-cold full allocation context hash 1093248920606587996: 10
 ; TOTALSIZESSINGLE: Total size for full allocation context hash 6792096022461663180 and single alloc type notcold: 10
 ; REMARKSINGLE: remark: memprof.cc:25:13: call in function main marked with memprof allocation attribute notcold
 ; TOTALSIZESSINGLE: Total size for full allocation context hash 15737101490731057601 and single alloc type cold: 10

@snehasish
Copy link

snehasish commented Dec 23, 2025

Is it possible to add the allocation type to contexttotalsize struct instead of creating a new pair?

@teresajohnson
Copy link
Contributor Author

teresajohnson commented Dec 23, 2025

Is it possible to add the allocation type to contexttotalsize struct instead of creating a new pair?

I didn't for a couple reasons. For one, the ContextTotalSize data structure is used other places, such as during the LTO step. And secondly, due to context trimming a summarized context with its allocation type can correspond to multiple context/size pairs. For the Trie during matching, the redundancy (in the new vector of pairs) is ok from a convenience standpoint, but during LTO I don't want to carry around extra memory.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@snehasish snehasish Awaiting requested review from snehasish

@kazutakahirata kazutakahirata Awaiting requested review from kazutakahirata

Assignees

No one assigned

Labels

llvm:analysis Includes value tracking, cost tables and constant folding llvm:transforms PGO Profile Guided Optimizations

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@teresajohnson @llvmbot @snehasish